{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '3'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !cp /home/husein/pure-text/ontonotes/augmentation-*-ontonotes5.json .\n",
    "# !cp /home/husein/translation/ontonotes5-train-test.json ."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "config.json\t\t\t       model.ckpt-400000.meta\r\n",
      "model.ckpt-400000.data-00000-of-00001  sp10m.cased.v10.model\r\n",
      "model.ckpt-400000.index\t\t       sp10m.cased.v10.vocab\r\n"
     ]
    }
   ],
   "source": [
    "# !wget https://f000.backblazeb2.com/file/malaya-model/bert-bahasa/albert-base-2020-04-10.tar.gz\n",
    "# !tar -zxf albert-base-2020-04-10.tar.gz\n",
    "!ls albert-base-2020-04-10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/husein/.local/lib/python3.6/site-packages/albert/lamb_optimizer.py:34: The name tf.train.Optimizer is deprecated. Please use tf.compat.v1.train.Optimizer instead.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "from albert import modeling\n",
    "from albert import optimization\n",
    "from albert import tokenization\n",
    "import tensorflow as tf\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/husein/.local/lib/python3.6/site-packages/albert/tokenization.py:240: The name tf.logging.info is deprecated. Please use tf.compat.v1.logging.info instead.\n",
      "\n",
      "INFO:tensorflow:loading sentence piece model\n"
     ]
    }
   ],
   "source": [
    "tokenizer = tokenization.FullTokenizer(\n",
    "      vocab_file='albert-base-2020-04-10/sp10m.cased.v10.vocab', do_lower_case=False,\n",
    "      spm_model_file='albert-base-2020-04-10/sp10m.cased.v10.model')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/husein/.local/lib/python3.6/site-packages/albert/modeling.py:116: The name tf.gfile.GFile is deprecated. Please use tf.io.gfile.GFile instead.\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<albert.modeling.AlbertConfig at 0x7f0b201bbac8>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "albert_config = modeling.AlbertConfig.from_json_file('albert-base-2020-04-10/config.json')\n",
    "albert_config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "prefix = 'https://f000.backblazeb2.com/file/malay-dataset/tagging/ontonotes5/'\n",
    "\n",
    "urls = \"\"\"\n",
    "augmentation-address-ontonotes5.json\n",
    "augmentation-event-ontonotes5.json\n",
    "augmentation-fac-ontonotes5.json\n",
    "augmentation-gpe-ontonotes5.json\n",
    "augmentation-language-ontonotes5.json\n",
    "augmentation-law-ontonotes5.json\n",
    "augmentation-loc-ontonotes5.json\n",
    "augmentation-norp-ontonotes5.json\n",
    "augmentation-org-ontonotes5.json\n",
    "augmentation-person-ontonotes5.json\n",
    "augmentation-product-ontonotes5.json\n",
    "augmentation-work-of-art-ontonotes5.json\n",
    "ontonotes5-train-test.json\n",
    "\"\"\"\n",
    "urls = list(filter(None, urls.split('\\n')))\n",
    "\n",
    "import os\n",
    "\n",
    "# uncomment to download\n",
    "# for url in urls:\n",
    "#     print(url)\n",
    "#     os.system(f'wget {prefix}{url}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dict_keys(['train_X', 'train_Y', 'test_X', 'test_Y'])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import json\n",
    "\n",
    "with open('ontonotes5-train-test.json') as fopen:\n",
    "    data = json.load(fopen)\n",
    "data.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_X = data['train_X']\n",
    "test_X = data['test_X']\n",
    "train_Y = data['train_Y']\n",
    "test_Y = data['test_Y']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "d = [\n",
    "    {'Tag': 'OTHER', 'Description': 'other'},\n",
    "    {'Tag': 'ADDRESS', 'Description': 'Address of physical location.'},\n",
    "    {'Tag': 'PERSON', 'Description': 'People, including fictional.'},\n",
    "    {\n",
    "        'Tag': 'NORP',\n",
    "        'Description': 'Nationalities or religious or political groups.',\n",
    "    },\n",
    "    {\n",
    "        'Tag': 'FAC',\n",
    "        'Description': 'Buildings, airports, highways, bridges, etc.',\n",
    "    },\n",
    "    {\n",
    "        'Tag': 'ORG',\n",
    "        'Description': 'Companies, agencies, institutions, etc.',\n",
    "    },\n",
    "    {'Tag': 'GPE', 'Description': 'Countries, cities, states.'},\n",
    "    {\n",
    "        'Tag': 'LOC',\n",
    "        'Description': 'Non-GPE locations, mountain ranges, bodies of water.',\n",
    "    },\n",
    "    {\n",
    "        'Tag': 'PRODUCT',\n",
    "        'Description': 'Objects, vehicles, foods, etc. (Not services.)',\n",
    "    },\n",
    "    {\n",
    "        'Tag': 'EVENT',\n",
    "        'Description': 'Named hurricanes, battles, wars, sports events, etc.',\n",
    "    },\n",
    "    {'Tag': 'WORK_OF_ART', 'Description': 'Titles of books, songs, etc.'},\n",
    "    {'Tag': 'LAW', 'Description': 'Named documents made into laws.'},\n",
    "    {'Tag': 'LANGUAGE', 'Description': 'Any named language.'},\n",
    "    {\n",
    "        'Tag': 'DATE',\n",
    "        'Description': 'Absolute or relative dates or periods.',\n",
    "    },\n",
    "    {'Tag': 'TIME', 'Description': 'Times smaller than a day.'},\n",
    "    {'Tag': 'PERCENT', 'Description': 'Percentage, including \"%\".'},\n",
    "    {'Tag': 'MONEY', 'Description': 'Monetary values, including unit.'},\n",
    "    {\n",
    "        'Tag': 'QUANTITY',\n",
    "        'Description': 'Measurements, as of weight or distance.',\n",
    "    },\n",
    "    {'Tag': 'ORDINAL', 'Description': '\"first\", \"second\", etc.'},\n",
    "    {\n",
    "        'Tag': 'CARDINAL',\n",
    "        'Description': 'Numerals that do not fall under another type.',\n",
    "    },\n",
    "]\n",
    "d = [d['Tag'] for d in d]\n",
    "d = ['PAD', 'X'] + d\n",
    "tag2idx = {i: no for no, i in enumerate(d)}\n",
    "idx2tag = {no: i for no, i in enumerate(d)}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "augmentation-org-ontonotes5.json\n",
      "35150 35150\n",
      "6300 6300\n",
      "augmentation-fac-ontonotes5.json\n",
      "17040 17040\n",
      "4260 4260\n",
      "augmentation-loc-ontonotes5.json\n",
      "13040 13040\n",
      "2460 2460\n",
      "augmentation-gpe-ontonotes5.json\n",
      "21060 21060\n",
      "0 0\n",
      "augmentation-work-of-art-ontonotes5.json\n",
      "4020 4020\n",
      "1020 1020\n",
      "augmentation-event-ontonotes5.json\n",
      "1665 1665\n",
      "0 0\n",
      "augmentation-person-ontonotes5.json\n",
      "37883 37883\n",
      "2530 2530\n",
      "augmentation-product-ontonotes5.json\n",
      "7040 7040\n",
      "1760 1760\n",
      "augmentation-law-ontonotes5.json\n",
      "12584 12584\n",
      "3161 3161\n",
      "augmentation-language-ontonotes5.json\n",
      "6860 6860\n",
      "0 0\n",
      "augmentation-norp-ontonotes5.json\n",
      "30660 30660\n",
      "5940 5940\n",
      "augmentation-address-ontonotes5.json\n",
      "57502 57502\n",
      "15106 15106\n"
     ]
    }
   ],
   "source": [
    "from glob import glob\n",
    "\n",
    "augmented = glob('augmentation-*-ontonotes5.json')\n",
    "\n",
    "for f in augmented:\n",
    "    print(f)\n",
    "    with open(f) as fopen:\n",
    "        data = json.load(fopen)\n",
    "        \n",
    "    print(len(data.get('train_X', [])), len(data.get('train_Y', [])))\n",
    "    print(len(data.get('test_X', [])), len(data.get('test_Y', [])))\n",
    "    \n",
    "    train_X.extend(data.get('train_X', []))\n",
    "    train_Y.extend(data.get('train_Y', []))\n",
    "    test_X.extend(data.get('test_X', []))\n",
    "    test_Y.extend(data.get('test_Y', []))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "\n",
    "def XY(strings):\n",
    "    left_train, right_train = strings[0], strings[1]\n",
    "    X, Y, MASK = [], [], []\n",
    "    for i in tqdm(range(len(left_train))):\n",
    "        left = [d for d in left_train[i]]\n",
    "        right = [d for d in right_train[i]]\n",
    "        bert_tokens = ['[CLS]']\n",
    "        y = ['PAD']\n",
    "        for no, orig_token in enumerate(left):\n",
    "            t = tokenizer.tokenize(orig_token)\n",
    "            bert_tokens.extend(t)\n",
    "            if len(t):\n",
    "                y.append(right[no])\n",
    "            y.extend(['X'] * (len(t) - 1))\n",
    "        bert_tokens.append('[SEP]')\n",
    "        y.append('PAD')\n",
    "        x = tokenizer.convert_tokens_to_ids(bert_tokens)\n",
    "        y = [tag2idx[i] for i in y]\n",
    "        input_mask = [1] * len(y)\n",
    "        if len(x) != len(y):\n",
    "            print(i)\n",
    "        X.append(x)\n",
    "        Y.append(y)\n",
    "        MASK.append(input_mask)\n",
    "    return [(X, Y, MASK,)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 40812/40812 [00:28<00:00, 1412.06it/s]\n",
      " 59%|█████▉    | 24232/40812 [00:17<00:12, 1325.70it/s]\n",
      "100%|██████████| 40812/40812 [00:29<00:00, 1390.98it/s]\n",
      "100%|██████████| 40812/40812 [00:28<00:00, 1407.71it/s]\n",
      " 59%|█████▉    | 24069/40812 [00:18<00:11, 1402.73it/s]\n",
      "100%|██████████| 40812/40812 [00:29<00:00, 1379.77it/s]\n",
      "100%|██████████| 40812/40812 [00:28<00:00, 1426.88it/s]\n",
      "100%|██████████| 40812/40812 [00:29<00:00, 1360.83it/s]\n",
      "100%|██████████| 40812/40812 [00:28<00:00, 1434.81it/s]\n",
      "100%|██████████| 40812/40812 [00:30<00:00, 1358.03it/s]\n",
      "100%|██████████| 40812/40812 [00:29<00:00, 1384.60it/s]\n",
      "100%|██████████| 40812/40812 [00:31<00:00, 1301.45it/s]\n",
      "100%|██████████| 40812/40812 [00:30<00:00, 1322.40it/s]\n",
      "100%|██████████| 40812/40812 [00:31<00:00, 1310.02it/s]\n",
      "100%|██████████| 40812/40812 [00:30<00:00, 1332.85it/s]\n",
      "100%|██████████| 40812/40812 [00:33<00:00, 1204.14it/s]\n",
      "100%|██████████| 40812/40812 [00:33<00:00, 1214.84it/s]\n"
     ]
    }
   ],
   "source": [
    "import cleaning\n",
    "t = cleaning.multiprocessing(train_X, train_Y, XY)\n",
    "train_X, train_Y, train_masks = [], [], []\n",
    "for x, y, m in t:\n",
    "    train_X.extend(x)\n",
    "    train_Y.extend(y)\n",
    "    train_masks.extend(m)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 5494/5494 [00:05<00:00, 1084.23it/s]\n",
      " 75%|███████▍  | 4094/5494 [00:04<00:01, 887.66it/s]]\n",
      "100%|██████████| 10/10 [00:00<00:00, 1060.59it/s]t/s]\n",
      "100%|██████████| 5494/5494 [00:05<00:00, 987.83it/s]]\n",
      "100%|██████████| 5494/5494 [00:05<00:00, 1020.95it/s]\n",
      "100%|██████████| 5494/5494 [00:05<00:00, 965.65it/s] \n",
      "100%|██████████| 5494/5494 [00:05<00:00, 983.17it/s]]\n",
      "100%|██████████| 5494/5494 [00:05<00:00, 930.61it/s]]\n",
      "100%|██████████| 5494/5494 [00:05<00:00, 927.69it/s]]\n",
      "100%|██████████| 5494/5494 [00:05<00:00, 982.87it/s]]\n",
      "100%|██████████| 5494/5494 [00:05<00:00, 1060.84it/s]\n",
      "100%|██████████| 5494/5494 [00:05<00:00, 971.18it/s] \n",
      "100%|██████████| 5494/5494 [00:05<00:00, 942.55it/s] \n",
      "100%|██████████| 5494/5494 [00:06<00:00, 907.00it/s]]\n",
      "100%|██████████| 5494/5494 [00:05<00:00, 982.23it/s]]\n",
      "100%|██████████| 5494/5494 [00:05<00:00, 1019.90it/s]\n",
      "100%|██████████| 5494/5494 [00:05<00:00, 946.49it/s] \n"
     ]
    }
   ],
   "source": [
    "t = cleaning.multiprocessing(test_X, test_Y, XY)\n",
    "test_X, test_Y, test_masks = [], [], []\n",
    "for x, y, m in t:\n",
    "    test_X.extend(x)\n",
    "    test_Y.extend(y)\n",
    "    test_masks.extend(m)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.utils import shuffle\n",
    "\n",
    "train_X, train_Y, train_masks = shuffle(train_X, train_Y, train_masks)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "BERT_INIT_CHKPNT = 'albert-base-2020-04-10/model.ckpt-400000'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "epoch = 5\n",
    "batch_size = 32\n",
    "warmup_proportion = 0.1\n",
    "num_train_steps = int(1000 / batch_size * epoch)\n",
    "num_warmup_steps = int(num_train_steps * warmup_proportion)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_initializer(initializer_range=0.02):\n",
    "    return tf.truncated_normal_initializer(stddev=initializer_range)\n",
    "\n",
    "class Model:\n",
    "    def __init__(\n",
    "        self,\n",
    "        dimension_output,\n",
    "        learning_rate = 2e-5,\n",
    "        training = True\n",
    "    ):\n",
    "        self.X = tf.placeholder(tf.int32, [None, None])\n",
    "        self.MASK = tf.placeholder(tf.int32, [None, None])\n",
    "        self.Y = tf.placeholder(tf.int32, [None, None])\n",
    "        self.maxlen = tf.shape(self.X)[1]\n",
    "        self.lengths = tf.count_nonzero(self.X, 1)\n",
    "        \n",
    "        model = modeling.AlbertModel(\n",
    "            config=albert_config,\n",
    "            is_training=training,\n",
    "            input_ids=self.X,\n",
    "            input_mask=self.MASK,\n",
    "            use_one_hot_embeddings=False)\n",
    "        output_layer = model.get_sequence_output()\n",
    "        output_layer = tf.layers.dense(\n",
    "            output_layer,\n",
    "            albert_config.hidden_size,\n",
    "            activation=tf.tanh,\n",
    "            kernel_initializer=create_initializer())\n",
    "        logits = tf.layers.dense(output_layer, dimension_output,\n",
    "                                         kernel_initializer=create_initializer())\n",
    "        y_t = self.Y\n",
    "        log_likelihood, transition_params = tf.contrib.crf.crf_log_likelihood(\n",
    "            logits, y_t, self.lengths\n",
    "        )\n",
    "        self.cost = tf.reduce_mean(-log_likelihood)\n",
    "        self.optimizer = tf.train.AdamOptimizer(\n",
    "            learning_rate = learning_rate\n",
    "        ).minimize(self.cost)\n",
    "        mask = tf.sequence_mask(self.lengths, maxlen = self.maxlen)\n",
    "        self.tags_seq, tags_score = tf.contrib.crf.crf_decode(\n",
    "            logits, transition_params, self.lengths\n",
    "        )\n",
    "        self.tags_seq = tf.identity(self.tags_seq, name = 'logits')\n",
    "\n",
    "        y_t = tf.cast(y_t, tf.int32)\n",
    "        self.prediction = tf.boolean_mask(self.tags_seq, mask)\n",
    "        mask_label = tf.boolean_mask(y_t, mask)\n",
    "        correct_pred = tf.equal(self.prediction, mask_label)\n",
    "        correct_index = tf.cast(correct_pred, tf.float32)\n",
    "        self.accuracy = tf.reduce_mean(tf.cast(correct_pred, tf.float32))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/husein/.local/lib/python3.6/site-packages/tensorflow_core/python/util/deprecation.py:507: calling count_nonzero (from tensorflow.python.ops.math_ops) with axis is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "reduction_indices is deprecated, use axis instead\n",
      "WARNING:tensorflow:From /home/husein/.local/lib/python3.6/site-packages/albert/modeling.py:194: The name tf.variable_scope is deprecated. Please use tf.compat.v1.variable_scope instead.\n",
      "\n",
      "WARNING:tensorflow:From /home/husein/.local/lib/python3.6/site-packages/albert/modeling.py:507: The name tf.get_variable is deprecated. Please use tf.compat.v1.get_variable instead.\n",
      "\n",
      "WARNING:tensorflow:From /home/husein/.local/lib/python3.6/site-packages/albert/modeling.py:588: The name tf.assert_less_equal is deprecated. Please use tf.compat.v1.assert_less_equal instead.\n",
      "\n",
      "WARNING:tensorflow:From /home/husein/.local/lib/python3.6/site-packages/albert/modeling.py:1025: The name tf.AUTO_REUSE is deprecated. Please use tf.compat.v1.AUTO_REUSE instead.\n",
      "\n",
      "WARNING:tensorflow:From /home/husein/.local/lib/python3.6/site-packages/albert/modeling.py:253: dense (from tensorflow.python.layers.core) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use keras.layers.Dense instead.\n",
      "WARNING:tensorflow:From /home/husein/.local/lib/python3.6/site-packages/tensorflow_core/python/layers/core.py:187: Layer.apply (from tensorflow.python.keras.engine.base_layer) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Please use `layer.__call__` method instead.\n",
      "WARNING:tensorflow:\n",
      "The TensorFlow contrib module will not be included in TensorFlow 2.0.\n",
      "For more information, please see:\n",
      "  * https://github.com/tensorflow/community/blob/master/rfcs/20180907-contrib-sunset.md\n",
      "  * https://github.com/tensorflow/addons\n",
      "  * https://github.com/tensorflow/io (for I/O related ops)\n",
      "If you depend on functionality not listed there, please file an issue.\n",
      "\n",
      "WARNING:tensorflow:From /home/husein/.local/lib/python3.6/site-packages/tensorflow_core/contrib/crf/python/ops/crf.py:99: where (from tensorflow.python.ops.array_ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use tf.where in 2.0, which has the same broadcast rule as np.where\n",
      "WARNING:tensorflow:From /home/husein/.local/lib/python3.6/site-packages/tensorflow_core/contrib/crf/python/ops/crf.py:213: dynamic_rnn (from tensorflow.python.ops.rnn) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Please use `keras.layers.RNN(cell)`, which is equivalent to this API\n",
      "INFO:tensorflow:Restoring parameters from albert-base-2020-04-10/model.ckpt-400000\n"
     ]
    }
   ],
   "source": [
    "dimension_output = len(tag2idx)\n",
    "learning_rate = 2e-5\n",
    "\n",
    "tf.reset_default_graph()\n",
    "sess = tf.InteractiveSession()\n",
    "model = Model(\n",
    "    dimension_output,\n",
    "    learning_rate\n",
    ")\n",
    "\n",
    "sess.run(tf.global_variables_initializer())\n",
    "var_lists = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope = 'bert')\n",
    "saver = tf.train.Saver(var_list = var_lists)\n",
    "saver.restore(sess, BERT_INIT_CHKPNT)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "def merge_sentencepiece_tokens_tagging(x, y):\n",
    "    new_paired_tokens = []\n",
    "    n_tokens = len(x)\n",
    "    rejected = ['[CLS]', '[SEP]']\n",
    "\n",
    "    i = 0\n",
    "\n",
    "    while i < n_tokens:\n",
    "\n",
    "        current_token, current_label = x[i], y[i]\n",
    "        if not current_token.startswith('▁') and current_token not in rejected:\n",
    "            previous_token, previous_label = new_paired_tokens.pop()\n",
    "            merged_token = previous_token\n",
    "            merged_label = [previous_label]\n",
    "            while (\n",
    "                not current_token.startswith('▁')\n",
    "                and current_token not in rejected\n",
    "            ):\n",
    "                merged_token = merged_token + current_token.replace('▁', '')\n",
    "                merged_label.append(current_label)\n",
    "                i = i + 1\n",
    "                current_token, current_label = x[i], y[i]\n",
    "            merged_label = merged_label[0]\n",
    "            new_paired_tokens.append((merged_token, merged_label))\n",
    "\n",
    "        else:\n",
    "            new_paired_tokens.append((current_token, current_label))\n",
    "            i = i + 1\n",
    "\n",
    "    words = [\n",
    "        i[0].replace('▁', '')\n",
    "        for i in new_paired_tokens\n",
    "        if i[0] not in rejected\n",
    "    ]\n",
    "    labels = [i[1] for i in new_paired_tokens if i[0] not in rejected]\n",
    "    return words, labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "string = 'KUALA LUMPUR : Sempena sambutan Aidilfitri minggu depan, Perdana Menteri Tun Dr Mahathir Mohamad dan Menteri Pengangkutan Anthony Loke Siew Fook menitipkan pesanan khas kepada orang ramai yang mahu pulang ke kampung halaman masing-masing. Dalam video pendek terbitan Jabatan Keselamatan Jalan Raya (JKJR) itu, Dr Mahathir menasihati mereka supaya berhenti berehat dan tidur sebentar  sekiranya mengantuk ketika memandu.'\n",
    "\n",
    "import re\n",
    "\n",
    "def entities_textcleaning(string, lowering = False):\n",
    "    \"\"\"\n",
    "    use by entities recognition, pos recognition and dependency parsing\n",
    "    \"\"\"\n",
    "    string = re.sub(r'[ ]+', ' ', string).strip()\n",
    "    original_string = string.split()\n",
    "    if lowering:\n",
    "        string = string.lower()\n",
    "    string = [\n",
    "        (original_string[no], word.title() if word.isupper() else word)\n",
    "        for no, word in enumerate(string.split())\n",
    "        if len(word)\n",
    "    ]\n",
    "    return [s[0] for s in string], [s[1] for s in string]\n",
    "\n",
    "def parse_X(left):\n",
    "    bert_tokens = ['[CLS]']\n",
    "    for no, orig_token in enumerate(left):\n",
    "        t = tokenizer.tokenize(orig_token)\n",
    "        bert_tokens.extend(t)\n",
    "    bert_tokens.append(\"[SEP]\")\n",
    "    input_mask = [1] * len(bert_tokens)\n",
    "    return tokenizer.convert_tokens_to_ids(bert_tokens), bert_tokens, input_mask\n",
    "\n",
    "sequence = entities_textcleaning(string)[1]\n",
    "parsed_sequence, bert_sequence, input_mask = parse_X(sequence)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[('Kuala', 'X'),\n",
       " ('Lumpur', 'DATE'),\n",
       " (':', 'TIME'),\n",
       " ('Sempena', 'LOC'),\n",
       " ('sambutan', 'LAW'),\n",
       " ('Aidilfitri', 'ORG'),\n",
       " ('minggu', 'OTHER'),\n",
       " ('depan,', 'QUANTITY'),\n",
       " ('Perdana', 'DATE'),\n",
       " ('Menteri', 'ORG'),\n",
       " ('Tun', 'TIME'),\n",
       " ('Dr', 'LOC'),\n",
       " ('Mahathir', 'QUANTITY'),\n",
       " ('Mohamad', 'ORG'),\n",
       " ('dan', 'TIME'),\n",
       " ('Menteri', 'TIME'),\n",
       " ('Pengangkutan', 'LAW'),\n",
       " ('Anthony', 'ORG'),\n",
       " ('Loke', 'TIME'),\n",
       " ('Siew', 'LAW'),\n",
       " ('Fook', 'ORDINAL'),\n",
       " ('menitipkan', 'ORG'),\n",
       " ('pesanan', 'MONEY'),\n",
       " ('khas', 'PRODUCT'),\n",
       " ('kepada', 'ORDINAL'),\n",
       " ('orang', 'ORG'),\n",
       " ('ramai', 'LOC'),\n",
       " ('yang', 'ADDRESS'),\n",
       " ('mahu', 'EVENT'),\n",
       " ('pulang', 'LOC'),\n",
       " ('ke', 'QUANTITY'),\n",
       " ('kampung', 'ORG'),\n",
       " ('halaman', 'TIME'),\n",
       " ('masing-masing.', 'LOC'),\n",
       " ('Dalam', 'WORK_OF_ART'),\n",
       " ('video', 'NORP'),\n",
       " ('pendek', 'WORK_OF_ART'),\n",
       " ('terbitan', 'MONEY'),\n",
       " ('Jabatan', 'MONEY'),\n",
       " ('Keselamatan', 'PRODUCT'),\n",
       " ('Jalan', 'LOC'),\n",
       " ('Raya', 'ORG'),\n",
       " ('(Jkjr)', 'X'),\n",
       " ('itu,', 'WORK_OF_ART'),\n",
       " ('Dr', 'NORP'),\n",
       " ('Mahathir', 'QUANTITY'),\n",
       " ('menasihati', 'ORG'),\n",
       " ('mereka', 'DATE'),\n",
       " ('supaya', 'TIME'),\n",
       " ('berhenti', 'ORDINAL'),\n",
       " ('berehat', 'ORG'),\n",
       " ('dan', 'X'),\n",
       " ('tidur', 'LAW'),\n",
       " ('sebentar', 'PERSON'),\n",
       " ('sekiranya', 'ORDINAL'),\n",
       " ('mengantuk', 'ORG'),\n",
       " ('ketika', 'OTHER'),\n",
       " ('memandu.', 'DATE')]"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "predicted = sess.run(model.tags_seq,\n",
    "                feed_dict = {\n",
    "                    model.X: [parsed_sequence],\n",
    "                    model.MASK: [input_mask]\n",
    "                },\n",
    "        )[0]\n",
    "merged = merge_sentencepiece_tokens_tagging(bert_sequence, [idx2tag[d] for d in predicted])\n",
    "list(zip(merged[0], merged[1]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "pad_sequences = tf.keras.preprocessing.sequence.pad_sequences"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 20407/20407 [2:00:46<00:00,  2.82it/s, accuracy=1, cost=1.38]       \n",
      "test minibatch loop: 100%|██████████| 2748/2748 [07:51<00:00,  5.83it/s, accuracy=0.997, cost=1.32] \n",
      "train minibatch loop:   0%|          | 0/20407 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "time taken: 7717.824207544327\n",
      "epoch: 0, training loss: 3.420557, training acc: 0.985193, valid loss: 4.921612, valid acc: 0.983174\n",
      "\n",
      "[('Kuala', 'GPE'), ('Lumpur', 'GPE'), (':', 'OTHER'), ('Sempena', 'OTHER'), ('sambutan', 'OTHER'), ('Aidilfitri', 'OTHER'), ('minggu', 'DATE'), ('depan,', 'DATE'), ('Perdana', 'ORG'), ('Menteri', 'ORG'), ('Tun', 'ORG'), ('Dr', 'ORG'), ('Mahathir', 'ORG'), ('Mohamad', 'ORG'), ('dan', 'OTHER'), ('Menteri', 'OTHER'), ('Pengangkutan', 'OTHER'), ('Anthony', 'PERSON'), ('Loke', 'PERSON'), ('Siew', 'PERSON'), ('Fook', 'PERSON'), ('menitipkan', 'OTHER'), ('pesanan', 'OTHER'), ('khas', 'OTHER'), ('kepada', 'OTHER'), ('orang', 'OTHER'), ('ramai', 'OTHER'), ('yang', 'OTHER'), ('mahu', 'OTHER'), ('pulang', 'OTHER'), ('ke', 'OTHER'), ('kampung', 'OTHER'), ('halaman', 'OTHER'), ('masing-masing.', 'OTHER'), ('Dalam', 'OTHER'), ('video', 'OTHER'), ('pendek', 'OTHER'), ('terbitan', 'OTHER'), ('Jabatan', 'ORG'), ('Keselamatan', 'ORG'), ('Jalan', 'ORG'), ('Raya', 'ORG'), ('(Jkjr)', 'ORG'), ('itu,', 'OTHER'), ('Dr', 'PERSON'), ('Mahathir', 'PERSON'), ('menasihati', 'OTHER'), ('mereka', 'OTHER'), ('supaya', 'OTHER'), ('berhenti', 'OTHER'), ('berehat', 'OTHER'), ('dan', 'OTHER'), ('tidur', 'OTHER'), ('sebentar', 'OTHER'), ('sekiranya', 'OTHER'), ('mengantuk', 'OTHER'), ('ketika', 'OTHER'), ('memandu.', 'OTHER')]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 20407/20407 [2:00:46<00:00,  2.82it/s, accuracy=1, cost=0.0455]     \n",
      "test minibatch loop: 100%|██████████| 2748/2748 [07:43<00:00,  5.93it/s, accuracy=1, cost=0.129]     \n",
      "train minibatch loop:   0%|          | 0/20407 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "time taken: 7710.071366548538\n",
      "epoch: 1, training loss: 0.625227, training acc: 0.997115, valid loss: 5.663307, valid acc: 0.984243\n",
      "\n",
      "[('Kuala', 'ORG'), ('Lumpur', 'ORG'), (':', 'OTHER'), ('Sempena', 'OTHER'), ('sambutan', 'OTHER'), ('Aidilfitri', 'DATE'), ('minggu', 'DATE'), ('depan,', 'OTHER'), ('Perdana', 'OTHER'), ('Menteri', 'OTHER'), ('Tun', 'ORG'), ('Dr', 'ORG'), ('Mahathir', 'ORG'), ('Mohamad', 'ORG'), ('dan', 'OTHER'), ('Menteri', 'OTHER'), ('Pengangkutan', 'OTHER'), ('Anthony', 'PERSON'), ('Loke', 'PERSON'), ('Siew', 'PERSON'), ('Fook', 'PERSON'), ('menitipkan', 'OTHER'), ('pesanan', 'OTHER'), ('khas', 'OTHER'), ('kepada', 'OTHER'), ('orang', 'OTHER'), ('ramai', 'OTHER'), ('yang', 'OTHER'), ('mahu', 'OTHER'), ('pulang', 'OTHER'), ('ke', 'OTHER'), ('kampung', 'OTHER'), ('halaman', 'OTHER'), ('masing-masing.', 'OTHER'), ('Dalam', 'OTHER'), ('video', 'OTHER'), ('pendek', 'OTHER'), ('terbitan', 'OTHER'), ('Jabatan', 'ORG'), ('Keselamatan', 'ORG'), ('Jalan', 'ORG'), ('Raya', 'ORG'), ('(Jkjr)', 'ORG'), ('itu,', 'OTHER'), ('Dr', 'PERSON'), ('Mahathir', 'PERSON'), ('menasihati', 'OTHER'), ('mereka', 'OTHER'), ('supaya', 'OTHER'), ('berhenti', 'OTHER'), ('berehat', 'OTHER'), ('dan', 'OTHER'), ('tidur', 'OTHER'), ('sebentar', 'OTHER'), ('sekiranya', 'OTHER'), ('mengantuk', 'OTHER'), ('ketika', 'OTHER'), ('memandu.', 'OTHER')]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop:   2%|▏         | 366/20407 [02:08<1:57:20,  2.85it/s, accuracy=0.995, cost=0.873] \n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-24-08fdd4304cbb>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     21\u001b[0m                 \u001b[0mmodel\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mX\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mbatch_x\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     22\u001b[0m                 \u001b[0mmodel\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mY\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mbatch_y\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 23\u001b[0;31m                 \u001b[0mmodel\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mMASK\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mbatch_masks\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     24\u001b[0m             },\n\u001b[1;32m     25\u001b[0m         )\n",
      "\u001b[0;32m~/.local/lib/python3.6/site-packages/tensorflow_core/python/client/session.py\u001b[0m in \u001b[0;36mrun\u001b[0;34m(self, fetches, feed_dict, options, run_metadata)\u001b[0m\n\u001b[1;32m    954\u001b[0m     \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    955\u001b[0m       result = self._run(None, fetches, feed_dict, options_ptr,\n\u001b[0;32m--> 956\u001b[0;31m                          run_metadata_ptr)\n\u001b[0m\u001b[1;32m    957\u001b[0m       \u001b[0;32mif\u001b[0m \u001b[0mrun_metadata\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    958\u001b[0m         \u001b[0mproto_data\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtf_session\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mTF_GetBuffer\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mrun_metadata_ptr\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.6/site-packages/tensorflow_core/python/client/session.py\u001b[0m in \u001b[0;36m_run\u001b[0;34m(self, handle, fetches, feed_dict, options, run_metadata)\u001b[0m\n\u001b[1;32m   1178\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mfinal_fetches\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0mfinal_targets\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mhandle\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0mfeed_dict_tensor\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1179\u001b[0m       results = self._do_run(handle, final_targets, final_fetches,\n\u001b[0;32m-> 1180\u001b[0;31m                              feed_dict_tensor, options, run_metadata)\n\u001b[0m\u001b[1;32m   1181\u001b[0m     \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1182\u001b[0m       \u001b[0mresults\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.6/site-packages/tensorflow_core/python/client/session.py\u001b[0m in \u001b[0;36m_do_run\u001b[0;34m(self, handle, target_list, fetch_list, feed_dict, options, run_metadata)\u001b[0m\n\u001b[1;32m   1357\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mhandle\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1358\u001b[0m       return self._do_call(_run_fn, feeds, fetches, targets, options,\n\u001b[0;32m-> 1359\u001b[0;31m                            run_metadata)\n\u001b[0m\u001b[1;32m   1360\u001b[0m     \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1361\u001b[0m       \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_do_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0m_prun_fn\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhandle\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfeeds\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfetches\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.6/site-packages/tensorflow_core/python/client/session.py\u001b[0m in \u001b[0;36m_do_call\u001b[0;34m(self, fn, *args)\u001b[0m\n\u001b[1;32m   1363\u001b[0m   \u001b[0;32mdef\u001b[0m \u001b[0m_do_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfn\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1364\u001b[0m     \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1365\u001b[0;31m       \u001b[0;32mreturn\u001b[0m \u001b[0mfn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1366\u001b[0m     \u001b[0;32mexcept\u001b[0m \u001b[0merrors\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mOpError\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1367\u001b[0m       \u001b[0mmessage\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcompat\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mas_text\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0me\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmessage\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.6/site-packages/tensorflow_core/python/client/session.py\u001b[0m in \u001b[0;36m_run_fn\u001b[0;34m(feed_dict, fetch_list, target_list, options, run_metadata)\u001b[0m\n\u001b[1;32m   1348\u001b[0m       \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_extend_graph\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1349\u001b[0m       return self._call_tf_sessionrun(options, feed_dict, fetch_list,\n\u001b[0;32m-> 1350\u001b[0;31m                                       target_list, run_metadata)\n\u001b[0m\u001b[1;32m   1351\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1352\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_prun_fn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mhandle\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfeed_dict\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfetch_list\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.6/site-packages/tensorflow_core/python/client/session.py\u001b[0m in \u001b[0;36m_call_tf_sessionrun\u001b[0;34m(self, options, feed_dict, fetch_list, target_list, run_metadata)\u001b[0m\n\u001b[1;32m   1441\u001b[0m     return tf_session.TF_SessionRun_wrapper(self._session, options, feed_dict,\n\u001b[1;32m   1442\u001b[0m                                             \u001b[0mfetch_list\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtarget_list\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1443\u001b[0;31m                                             run_metadata)\n\u001b[0m\u001b[1;32m   1444\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1445\u001b[0m   \u001b[0;32mdef\u001b[0m \u001b[0m_call_tf_sessionprun\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhandle\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfeed_dict\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfetch_list\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "import time\n",
    "\n",
    "for e in range(epoch):\n",
    "    lasttime = time.time()\n",
    "    train_acc, train_loss, test_acc, test_loss = [], [], [], []\n",
    "    pbar = tqdm(\n",
    "        range(0, len(train_X), batch_size), desc = 'train minibatch loop'\n",
    "    )\n",
    "    for i in pbar:\n",
    "        index = min(i + batch_size, len(train_X))\n",
    "        batch_x = train_X[i : index]\n",
    "        batch_y = train_Y[i : index]\n",
    "        batch_masks = train_masks[i : index]\n",
    "        batch_x = pad_sequences(batch_x, padding='post')\n",
    "        batch_y = pad_sequences(batch_y, padding='post')\n",
    "        batch_masks = pad_sequences(batch_masks, padding='post')\n",
    "        \n",
    "        acc, cost, _ = sess.run(\n",
    "            [model.accuracy, model.cost, model.optimizer],\n",
    "            feed_dict = {\n",
    "                model.X: batch_x,\n",
    "                model.Y: batch_y,\n",
    "                model.MASK: batch_masks,\n",
    "            },\n",
    "        )\n",
    "        assert not np.isnan(cost)\n",
    "        train_loss.append(cost)\n",
    "        train_acc.append(acc)\n",
    "        pbar.set_postfix(cost = cost, accuracy = acc)\n",
    "    \n",
    "    pbar = tqdm(\n",
    "        range(0, len(test_X), batch_size), desc = 'test minibatch loop'\n",
    "    )\n",
    "    for i in pbar:\n",
    "        index = min(i + batch_size, len(test_X))\n",
    "        batch_x = test_X[i : index]\n",
    "        batch_y = test_Y[i : index]\n",
    "        batch_masks = test_masks[i : index]\n",
    "        batch_x = pad_sequences(batch_x, padding='post')\n",
    "        batch_y = pad_sequences(batch_y, padding='post')\n",
    "        batch_masks = pad_sequences(batch_masks, padding='post')\n",
    "        \n",
    "        acc, cost = sess.run(\n",
    "            [model.accuracy, model.cost],\n",
    "            feed_dict = {\n",
    "                model.X: batch_x,\n",
    "                model.Y: batch_y,\n",
    "                model.MASK: batch_masks,\n",
    "            },\n",
    "        )\n",
    "        assert not np.isnan(cost)\n",
    "        test_loss.append(cost)\n",
    "        test_acc.append(acc)\n",
    "        pbar.set_postfix(cost = cost, accuracy = acc)\n",
    "    \n",
    "    train_loss = np.mean(train_loss)\n",
    "    train_acc = np.mean(train_acc)\n",
    "    test_loss = np.mean(test_loss)\n",
    "    test_acc = np.mean(test_acc)\n",
    "\n",
    "    print('time taken:', time.time() - lasttime)\n",
    "    print(\n",
    "        'epoch: %d, training loss: %f, training acc: %f, valid loss: %f, valid acc: %f\\n'\n",
    "        % (e, train_loss, train_acc, test_loss, test_acc)\n",
    "    )\n",
    "    predicted = sess.run(model.tags_seq,\n",
    "                feed_dict = {\n",
    "                    model.X: [parsed_sequence],\n",
    "                    model.MASK: [input_mask]\n",
    "                },\n",
    "        )[0]\n",
    "    merged = merge_sentencepiece_tokens_tagging(bert_sequence, [idx2tag[d] for d in predicted])\n",
    "    print(list(zip(merged[0], merged[1])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'albert-base-entities/model.ckpt'"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "saver = tf.train.Saver(tf.trainable_variables())\n",
    "saver.save(sess, 'albert-base-entities/model.ckpt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/husein/.local/lib/python3.6/site-packages/tensorflow_core/python/client/session.py:1750: UserWarning: An interactive session is already active. This can cause out-of-memory errors in some cases. You must explicitly call `InteractiveSession.close()` to release resources held by the other session(s).\n",
      "  warnings.warn('An interactive session is already active. This can '\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Restoring parameters from albert-base-entities/model.ckpt\n"
     ]
    }
   ],
   "source": [
    "dimension_output = len(tag2idx)\n",
    "learning_rate = 2e-5\n",
    "\n",
    "tf.reset_default_graph()\n",
    "sess = tf.InteractiveSession()\n",
    "model = Model(\n",
    "    dimension_output,\n",
    "    learning_rate,\n",
    "    training = False\n",
    ")\n",
    "\n",
    "sess.run(tf.global_variables_initializer())\n",
    "saver = tf.train.Saver(tf.trainable_variables())\n",
    "saver.restore(sess, 'albert-base-entities/model.ckpt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pred2label(pred):\n",
    "    out = []\n",
    "    for pred_i in pred:\n",
    "        out_i = []\n",
    "        for p in pred_i:\n",
    "            out_i.append(idx2tag[p])\n",
    "        out.append(out_i)\n",
    "    return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "validation minibatch loop: 100%|██████████| 2748/2748 [07:51<00:00,  5.83it/s]\n"
     ]
    }
   ],
   "source": [
    "real_Y, predict_Y = [], []\n",
    "\n",
    "pbar = tqdm(\n",
    "    range(0, len(test_X), batch_size), desc = 'validation minibatch loop'\n",
    ")\n",
    "for i in pbar:\n",
    "    index = min(i + batch_size, len(test_X))\n",
    "    batch_x = test_X[i : index]\n",
    "    batch_y = test_Y[i : index]\n",
    "    batch_masks = test_masks[i : index]\n",
    "    batch_x = pad_sequences(batch_x, padding='post')\n",
    "    batch_y = pad_sequences(batch_y, padding='post')\n",
    "    batch_masks = pad_sequences(batch_masks, padding='post')\n",
    "    predicted = pred2label(sess.run(model.tags_seq,\n",
    "            feed_dict = {\n",
    "                model.X: batch_x,\n",
    "                model.MASK: batch_masks,\n",
    "            },\n",
    "    ))\n",
    "    real = pred2label(batch_y)\n",
    "    predict_Y.extend(predicted)\n",
    "    real_Y.extend(real)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "temp_real_Y = []\n",
    "for r in real_Y:\n",
    "    temp_real_Y.extend(r)\n",
    "    \n",
    "temp_predict_Y = []\n",
    "for r in predict_Y:\n",
    "    temp_predict_Y.extend(r)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "              precision    recall  f1-score   support\n",
      "\n",
      "     ADDRESS    0.99832   0.99969   0.99901     93446\n",
      "    CARDINAL    0.93291   0.89046   0.91119     48255\n",
      "        DATE    0.94941   0.93032   0.93977    126548\n",
      "       EVENT    0.89330   0.92506   0.90890      5711\n",
      "         FAC    0.92540   0.91257   0.91894     27392\n",
      "         GPE    0.93484   0.93404   0.93444    101357\n",
      "    LANGUAGE    0.87207   0.92528   0.89789       803\n",
      "         LAW    0.95567   0.95140   0.95353     24834\n",
      "         LOC    0.91754   0.92362   0.92057     34538\n",
      "       MONEY    0.85349   0.87696   0.86507     30032\n",
      "        NORP    0.91698   0.91416   0.91557     57014\n",
      "     ORDINAL    0.89159   0.93320   0.91192      6213\n",
      "         ORG    0.95070   0.92537   0.93786    219533\n",
      "       OTHER    0.99000   0.99315   0.99157   3553350\n",
      "         PAD    1.00000   1.00000   1.00000   1291289\n",
      "     PERCENT    0.96746   0.95953   0.96348     21722\n",
      "      PERSON    0.93748   0.96710   0.95206    101981\n",
      "     PRODUCT    0.84749   0.77832   0.81143     11124\n",
      "    QUANTITY    0.92798   0.92079   0.92437     11614\n",
      "        TIME    0.91058   0.88413   0.89716      9502\n",
      " WORK_OF_ART    0.88983   0.77196   0.82671     13800\n",
      "           X    0.99917   0.99799   0.99858   1351758\n",
      "\n",
      "    accuracy                        0.98714   7141816\n",
      "   macro avg    0.93010   0.92341   0.92636   7141816\n",
      "weighted avg    0.98707   0.98714   0.98707   7141816\n",
      "\n"
     ]
    }
   ],
   "source": [
    "from sklearn.metrics import classification_report\n",
    "print(classification_report(temp_real_Y, temp_predict_Y, digits = 5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['Placeholder',\n",
       " 'Placeholder_1',\n",
       " 'Placeholder_2',\n",
       " 'bert/embeddings/word_embeddings',\n",
       " 'bert/embeddings/token_type_embeddings',\n",
       " 'bert/embeddings/position_embeddings',\n",
       " 'bert/embeddings/LayerNorm/gamma',\n",
       " 'bert/encoder/embedding_hidden_mapping_in/kernel',\n",
       " 'bert/encoder/embedding_hidden_mapping_in/bias',\n",
       " 'bert/encoder/transformer/group_0/inner_group_0/attention_1/self/query/kernel',\n",
       " 'bert/encoder/transformer/group_0/inner_group_0/attention_1/self/query/bias',\n",
       " 'bert/encoder/transformer/group_0/inner_group_0/attention_1/self/key/kernel',\n",
       " 'bert/encoder/transformer/group_0/inner_group_0/attention_1/self/key/bias',\n",
       " 'bert/encoder/transformer/group_0/inner_group_0/attention_1/self/value/kernel',\n",
       " 'bert/encoder/transformer/group_0/inner_group_0/attention_1/self/value/bias',\n",
       " 'bert/encoder/transformer/group_0/inner_group_0/attention_1/output/dense/kernel',\n",
       " 'bert/encoder/transformer/group_0/inner_group_0/attention_1/output/dense/bias',\n",
       " 'bert/encoder/transformer/group_0/inner_group_0/LayerNorm/gamma',\n",
       " 'bert/encoder/transformer/group_0/inner_group_0/ffn_1/intermediate/dense/kernel',\n",
       " 'bert/encoder/transformer/group_0/inner_group_0/ffn_1/intermediate/dense/bias',\n",
       " 'bert/encoder/transformer/group_0/inner_group_0/ffn_1/intermediate/output/dense/kernel',\n",
       " 'bert/encoder/transformer/group_0/inner_group_0/ffn_1/intermediate/output/dense/bias',\n",
       " 'bert/encoder/transformer/group_0/inner_group_0/LayerNorm_1/gamma',\n",
       " 'bert/pooler/dense/kernel',\n",
       " 'bert/pooler/dense/bias',\n",
       " 'dense/kernel',\n",
       " 'dense/bias',\n",
       " 'dense_1/kernel',\n",
       " 'dense_1/bias',\n",
       " 'transitions',\n",
       " 'logits']"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "strings = ','.join(\n",
    "    [\n",
    "        n.name\n",
    "        for n in tf.get_default_graph().as_graph_def().node\n",
    "        if ('Variable' in n.op\n",
    "        or 'Placeholder' in n.name\n",
    "        or 'logits' in n.name\n",
    "        or 'alphas' in n.name\n",
    "        or 'self/Softmax' in n.name)\n",
    "        and 'Adam' not in n.name\n",
    "        and 'beta' not in n.name\n",
    "        and 'global_step' not in n.name\n",
    "    ]\n",
    ")\n",
    "strings.split(',')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "def freeze_graph(model_dir, output_node_names):\n",
    "\n",
    "    if not tf.gfile.Exists(model_dir):\n",
    "        raise AssertionError(\n",
    "            \"Export directory doesn't exists. Please specify an export \"\n",
    "            'directory: %s' % model_dir\n",
    "        )\n",
    "\n",
    "    checkpoint = tf.train.get_checkpoint_state(model_dir)\n",
    "    input_checkpoint = checkpoint.model_checkpoint_path\n",
    "\n",
    "    absolute_model_dir = '/'.join(input_checkpoint.split('/')[:-1])\n",
    "    output_graph = absolute_model_dir + '/frozen_model.pb'\n",
    "    clear_devices = True\n",
    "    with tf.Session(graph = tf.Graph()) as sess:\n",
    "        saver = tf.train.import_meta_graph(\n",
    "            input_checkpoint + '.meta', clear_devices = clear_devices\n",
    "        )\n",
    "        saver.restore(sess, input_checkpoint)\n",
    "        output_graph_def = tf.graph_util.convert_variables_to_constants(\n",
    "            sess,\n",
    "            tf.get_default_graph().as_graph_def(),\n",
    "            output_node_names.split(','),\n",
    "        )\n",
    "        with tf.gfile.GFile(output_graph, 'wb') as f:\n",
    "            f.write(output_graph_def.SerializeToString())\n",
    "        print('%d ops in the final graph.' % len(output_graph_def.node))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Restoring parameters from albert-base-entities/model.ckpt\n",
      "WARNING:tensorflow:From <ipython-input-32-9a7215a4e58a>:23: convert_variables_to_constants (from tensorflow.python.framework.graph_util_impl) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use `tf.compat.v1.graph_util.convert_variables_to_constants`\n",
      "WARNING:tensorflow:From /home/husein/.local/lib/python3.6/site-packages/tensorflow_core/python/framework/graph_util_impl.py:277: extract_sub_graph (from tensorflow.python.framework.graph_util_impl) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use `tf.compat.v1.graph_util.extract_sub_graph`\n",
      "INFO:tensorflow:Froze 30 variables.\n",
      "INFO:tensorflow:Converted 30 variables to const ops.\n",
      "3539 ops in the final graph.\n"
     ]
    }
   ],
   "source": [
    "freeze_graph('albert-base-entities', strings)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/husein/.local/lib/python3.6/site-packages/tensorflow_core/python/client/session.py:1750: UserWarning: An interactive session is already active. This can cause out-of-memory errors in some cases. You must explicitly call `InteractiveSession.close()` to release resources held by the other session(s).\n",
      "  warnings.warn('An interactive session is already active. This can '\n"
     ]
    }
   ],
   "source": [
    "def load_graph(frozen_graph_filename):\n",
    "    with tf.gfile.GFile(frozen_graph_filename, 'rb') as f:\n",
    "        graph_def = tf.GraphDef()\n",
    "        graph_def.ParseFromString(f.read())\n",
    "    with tf.Graph().as_default() as graph:\n",
    "        tf.import_graph_def(graph_def)\n",
    "    return graph\n",
    "\n",
    "g = load_graph('albert-base-entities/frozen_model.pb')\n",
    "x = g.get_tensor_by_name('import/Placeholder:0')\n",
    "mask = g.get_tensor_by_name('import/Placeholder_1:0')\n",
    "logits = g.get_tensor_by_name('import/logits:0')\n",
    "test_sess = tf.InteractiveSession(graph = g)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "string = 'Kyrgios, 25 , membuat pesanan itu kerana menyedaari pelbagai kesukaran menimpa rakyat Australia ekoran perintah kawalan pergerakan yang diumumkan Mac lalu bagi memerangi wabak COVID-19 di negara berkenaan. Pemain tenis ranking ke-40 dunia yang dilahirkan di Canberra itu meminta pengikut dan penyokongnya agar jangan tidur dalam keadaan perut kosong dalam hantaran Instagram yang meraih lebih 92,000 tanda suka.'\n",
    "\n",
    "ori, sequence = entities_textcleaning(string)\n",
    "parsed_sequence, bert_sequence, input_mask = parse_X(sequence)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[('Kyrgios,', 'PERSON'), ('25', 'OTHER'), (',', 'OTHER'), ('membuat', 'OTHER'), ('pesanan', 'OTHER'), ('itu', 'OTHER'), ('kerana', 'OTHER'), ('menyedaari', 'OTHER'), ('pelbagai', 'OTHER'), ('kesukaran', 'OTHER'), ('menimpa', 'OTHER'), ('rakyat', 'OTHER'), ('Australia', 'NORP'), ('ekoran', 'OTHER'), ('perintah', 'OTHER'), ('kawalan', 'OTHER'), ('pergerakan', 'OTHER'), ('yang', 'OTHER'), ('diumumkan', 'OTHER'), ('Mac', 'DATE'), ('lalu', 'DATE'), ('bagi', 'OTHER'), ('memerangi', 'OTHER'), ('wabak', 'OTHER'), ('Covid-19', 'OTHER'), ('di', 'OTHER'), ('negara', 'OTHER'), ('berkenaan.', 'OTHER'), ('Pemain', 'OTHER'), ('tenis', 'OTHER'), ('ranking', 'OTHER'), ('ke-40', 'ORDINAL'), ('dunia', 'OTHER'), ('yang', 'OTHER'), ('dilahirkan', 'OTHER'), ('di', 'OTHER'), ('Canberra', 'GPE'), ('itu', 'OTHER'), ('meminta', 'OTHER'), ('pengikut', 'OTHER'), ('dan', 'OTHER'), ('penyokongnya', 'OTHER'), ('agar', 'OTHER'), ('jangan', 'OTHER'), ('tidur', 'OTHER'), ('dalam', 'OTHER'), ('keadaan', 'OTHER'), ('perut', 'OTHER'), ('kosong', 'OTHER'), ('dalam', 'OTHER'), ('hantaran', 'OTHER'), ('Instagram', 'ORG'), ('yang', 'OTHER'), ('meraih', 'OTHER'), ('lebih', 'CARDINAL'), ('92,000', 'CARDINAL'), ('tanda', 'OTHER'), ('suka.', 'OTHER')]\n"
     ]
    }
   ],
   "source": [
    "predicted = test_sess.run(logits,\n",
    "            feed_dict = {\n",
    "                x: [parsed_sequence],\n",
    "                mask: [input_mask]\n",
    "            },\n",
    "    )[0]\n",
    "merged = merge_sentencepiece_tokens_tagging(bert_sequence, [idx2tag[d] for d in predicted])\n",
    "print(list(zip(merged[0], merged[1])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "from tensorflow.tools.graph_transforms import TransformGraph\n",
    "tf.set_random_seed(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From <ipython-input-38-59d47e8570a3>:12: FastGFile.__init__ (from tensorflow.python.platform.gfile) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use tf.gfile.GFile.\n",
      "albert-base-entities/frozen_model.pb ['Placeholder', 'Placeholder_1']\n"
     ]
    }
   ],
   "source": [
    "transforms = ['add_default_attributes',\n",
    "             'remove_nodes(op=Identity, op=CheckNumerics, op=Dropout)',\n",
    "             'fold_batch_norms',\n",
    "             'fold_old_batch_norms',\n",
    "             'quantize_weights(fallback_min=-10, fallback_max=10)',\n",
    "             'strip_unused_nodes',\n",
    "             'sort_by_execution_order']\n",
    "\n",
    "pb = 'albert-base-entities/frozen_model.pb'\n",
    "\n",
    "input_graph_def = tf.GraphDef()\n",
    "with tf.gfile.FastGFile(pb, 'rb') as f:\n",
    "    input_graph_def.ParseFromString(f.read())\n",
    "\n",
    "inputs = ['Placeholder', 'Placeholder_1']\n",
    "outputs = ['dense/BiasAdd']\n",
    "\n",
    "print(pb, inputs)\n",
    "\n",
    "transformed_graph_def = TransformGraph(input_graph_def, \n",
    "                                       inputs,\n",
    "                                       ['logits'] + outputs, transforms)\n",
    "\n",
    "with tf.gfile.GFile(f'{pb}.quantized', 'wb') as f:\n",
    "    f.write(transformed_graph_def.SerializeToString())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
